import os

from dotenv import load_dotenv
from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
from langchain_core.messages.tool import tool_call
from langchain_core.prompts import PromptTemplate, ChatPromptTemplate, FewShotChatMessagePromptTemplate
from pydantic import BaseModel, Field
from typing import Optional

from tools import quote_detection, quote_extract


class Classification(BaseModel):
    """分类结果"""
    has_citation: str = Field(description="是否包含对经典原文的引用", enum=["是", "否"])
    citation: Optional[str] = Field(description="引用的文本范围")
    reason: Optional[str] = Field(description="判断是否包含引用的理由")


class Regex(BaseModel):
    """正则表达式检测和生成结果"""
    has_regex: str = Field(description="是否包含对文本范围的描述", enum=["是", "否"])
    regex: Optional[list[str]] = Field(description="生成的一个或多个正则表达式")
    reason: Optional[str] = Field(description="生成正则表达式的理由")


class AnnotationMatch(BaseModel):
    """注释匹配结果"""
    is_annotation: str = Field(description="古文1是否是对古文2中的注释/解释", enum=["是", "否"])
    key_words: Optional[str] = Field(description="匹配的关键词")
    reason: Optional[str] = Field(description="判断理由")


class Agents:
    """基础agent类，封装通用功能"""

    def __init__(self, temperature: float = 0.7):
        self.temperature = temperature
        self.agents = {}
        # 初始化模型连接等通用操作
        self.model = self._initialize_model()
        self.construct_agents()

    def _initialize_model(self):
        """初始化模型连接，子类可重写"""
        api_key = os.getenv("DeepSeek_API_KEY")
        if not api_key:
            # 加载当前目录的 .env 文件（文件内容类似：SECRET_KEY=xxx）
            load_dotenv()
            # 从加载的变量中获取值，用法同 os.getenv()
            api_key = os.getenv("DeepSeek_API_KEY")

        # https://python.langchain.com/api_reference/deepseek/chat_models/langchain_deepseek.chat_models.ChatDeepSeek.html
        from langchain_deepseek import ChatDeepSeek

        llm = ChatDeepSeek(
            model="deepseek-chat",
            temperature=self.temperature,
            max_tokens=1000,
            timeout=None,
            max_retries=2,
            api_key=api_key,
            # other params...
        )
        return llm

    def build_template(self, system_prompt, examples, use_tools=False, param_template=None):
        # 创建系统提示消息
        messages = [
            ("system", system_prompt),
        ]

        if use_tools:
            # 使用工具的情况下，直接添加示例消息
            messages.extend(examples)
        else:
            # 不使用工具的情况下，根据是否提供了param_template来决定使用哪种模板
            if param_template:
                # 使用自定义多参数模板
                example_prompt = ChatPromptTemplate.from_messages(
                    [
                        ("human", param_template),
                        ("ai", "{output}"),
                    ]
                )
            else:
                # 使用单参数模板
                example_prompt = ChatPromptTemplate.from_messages(
                    [
                        ("human", "{input}"),
                        ("ai", "{output}"),
                    ]
                )

            few_shot_prompt = FewShotChatMessagePromptTemplate(
                example_prompt=example_prompt,
                examples=examples,
            )
            messages.append(few_shot_prompt)

        # 最后添加用户输入提示
        if param_template:
            messages.append(("human", param_template))
        else:
            messages.append(("human", "{input}"))

        # 创建并返回统一的ChatPromptTemplate
        final_prompt = ChatPromptTemplate.from_messages(messages)
        return final_prompt

    def construct_regex_agent(self):
        system_prompt = "你是一个古文研究专家，并且精通正则表达式。请理解并判断下面文本中是否包含对文本范围的描述，如果包含，请根据文本生成正则表达式。注意：对文本范围的描述常使用\"從\"、\"訖\"等关键词，且与佛经对应内容通常以引号「」包裹，此外〈〉表示书名号，常用来标注章节题目表示文本范围，尤其是本章节开始位置。"

        # 定义参数模板和参数名称列表
        param_template = "原文：{text} 包含的章节名称：{chapter_names} 引号内容：{quote}"

        # 示例使用与参数模板匹配的结构
        examples = [
            {
                "text": "從「為諸菩薩說大乘經」訖「以佛舍利起七寶塔」。",
                "chapter_names": [],
                "quote": ["為諸菩薩說大乘經", "以佛舍利起七寶塔"],
                "output": '{"has_regex": "是", "regex": ["r\\[(為諸菩薩說大乘經).*?(以佛舍利起七寶塔)\\]"], "reason": "從是从哪里开始的意思，指出了引用的开始位置，訖是到哪里结束的意思，指出了引用的结束位置"}'
            },
            {
                "text": "《大智論》云：「外國云分陀利，此間言蓮花也。」",
                "chapter_names": [],
                "quote": ["「外國云分陀利，此間言蓮花也。」"],
                "output": '{"has_regex": "否", "reason": "此文本中没有明确的范围描述，无法生成正则表达式。"}'
            }
        ]

        # 传入参数模板和参数名称列表
        template = self.build_template(system_prompt, examples, use_tools=False, param_template=param_template)
        llm_with_tools = self.model.with_structured_output(Regex)
        agent = template | llm_with_tools
        return agent

    def construct_annotation_match_agent(self):
        system_prompt = "你是一个古文研究专家，你的任务是判断给定的古文1是否是对古文2中的注释/解释"
        # 将messages修改为examples列表，包含input和output键值对
        text1 = """「諸漏已盡」，此下是第四歎德，凡有五句，詺為依位歎德。初言「諸漏已盡無復煩惱」，此兩句歎羅漢不生之德。根本無明本是漏也，今所明者是業及潤生，明諸阿羅漢業及煩惱生因既盡，故言不生也。"""
        text2 = """一時，佛住王舍城耆闍崛山中，與大比丘眾萬二千人俱，皆是阿羅漢，諸漏已盡，無復煩惱，逮得己利，盡諸有結，心得自在。"""
        # example_prompt = PromptTemplate.from_template(f"古文1: {text1} 古文2: {text2} \n 回答：{output}")
        examples = [
            {
                "text1": text1,
                "text2": text2,
                "output": '{"is_annotation": "是", "key_words": "諸漏已盡", "reason": "引号中的諸漏已盡正是古文2中的关键词"}'
            },
        ]
        # 使用build_template创建模板（与construct_regex_agent保持一致）
        template = self.build_template(system_prompt, examples)
        # 绑定工具（保留原有的quote_detection工具）并设置结构化输出
        llm_with_tools = self.model.bind_tools([quote_detection, quote_extract]).with_structured_output(AnnotationMatch)
        # 通过模板和工具模型构建agent
        agent = template | llm_with_tools
        return agent

    def construct_judge_agent(self):
        system_prompt = "你是一个古文研究专家，你的任务是判断给定的古文中是否包含对佛经原文的引用。注意：对佛经原文的引用必须是直接引用，不能是对佛经原文的解释或注释。并且被引用部分常使用古代引号「」包裹。引用的部分可能包括佛经章节名称、佛经章节内容、佛经原文、佛经原文片段等。"
        # 将messages对话历史修改为examples示例列表，包含input和output键值对
        examples = [
            {
                "input": "原文：「韙哉主言！夫天台之道，迨至荊溪克任荷負，尤加光顯。竊閱其箋釋於玄文，台家宗體一翕一張，靡不昂然精覈震發。然至析丹丘㪺酌於光宅，鋤剗乎玄論之意，令人嘅焉不能無議。蓋不窺二書故也。其嘉祥《玄論》已行于世，此疏尚未及翻刻，予亦甚歉于衷。以誼不忍辭，丞欲鏤梓以公同好。」包含的章节名称：[] 引号内容：韙哉主言！夫天台之道，迨至荊溪克任荷負，尤加光顯。竊閱其箋釋於玄文，台家宗體一翕一張，靡不昂然精覈震發。然至析丹丘㪺酌於光宅，鋤剗乎玄論之意，令人嘅焉不能無議。蓋不窺二書故也。其嘉祥《玄論》已行于世，此疏尚未及翻刻，予亦甚歉于衷。以誼不忍辭，丞欲鏤梓以公同好。",
                "output": "{'has_citation': '否', 'reason': '虽然出现了引号，但引号中的内容并不是对佛经原文的引用。从引号内内容长度占比也可以看出，这不是对佛经原文的引用。并且此片段中没有包含任何章节名称。'}"
            },
            {
                "input": "原文：第三，從「說是普賢菩薩勸發品時」盡經，廣明時眾聞法受益多少，結成自行之義也。包含的章节名称：[] 引号内容：說是普賢菩薩勸發品時",
                "output": '{"has_citation": "是", "reason": "這個句子包含佛經原文的引用。引用了佛經原文中的「說是普賢菩薩勸發品時蓮華」，並對這個文本作出了解釋。 Although 這個片段中沒有包含任何章节名称。"}'
            },
            {
                "input": "從〈方便品〉初，訖「佛以方便力，示以三乘教，眾生處處著，引之令得出」以來，是第一，詺為略說開三顯一動執生疑。包含的章节名称：[方便品] 引号内容：佛以方便力，示以三乘教，眾生處處著，引之令得出",
                "output": '{"has_citation": "是", "reason": "這個句子包含佛經原文的引用。引用了佛經章節名稱「方便品」，以及佛經文本中的「佛以方便力，示以三乘教，眾生處處著，引之令得出」。", "citation": ["方便品", "佛以方便力，示以三乘教，眾生處處著，引之令得出"]}'
            }
        ]
        # 使用build_template创建提示模板（与construct_regex_agent保持一致）
        template = self.build_template(system_prompt, examples)
        # 绑定工具（保留原有的quote_detection工具）
        llm = self.model.bind_tools([quote_detection, quote_extract]).with_structured_output(Classification)
        # 通过模板和工具模型构建agent
        agent = template | llm
        return agent

    def construct_agents(self):
        self.agents["judge"] = self.construct_judge_agent()
        self.agents["regex"] = self.construct_regex_agent()
        self.agents["annotation_match"] = self.construct_annotation_match_agent()


'''
agent = Agent(model=model, tools=[calculator])
response = agent("What is 2+2")
print(response)
'''
